import collections

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.externals import joblib
import pylab as pl
import numpy as np

import stocknews
import stockdata
import datetime

# Settings
MAX_FEATURES = 200
COMPANY = 'GOOG'
START_DATE = "2013-11-01"
END_DATE = "2013-12-05"

CLF_FILE = 'data/classifiersX/Naive Bayes.pkl'
AVERAGE_SPAN_DAYS = 4

# Prepare
classifier = joblib.load(CLF_FILE)
#articles_train = stocknews.StockNews(db_file_name='data/_stocknews_db_639-stocks_2011-10-15-2013-10-15.shelve')
articles_test = stocknews.StockNews(db_file_name='data/_stocknews_db_639-stocks_2013-10-15-2013-11-31.shelve')
articles = stocknews.StockNews([COMPANY], START_DATE, END_DATE)
stock_prices = stockdata.StockData(COMPANY, START_DATE, END_DATE)
vectorizer = TfidfVectorizer(max_features=MAX_FEATURES, stop_words='english')

# Classify articles
X_predict = vectorizer.fit_transform([article for article in articles_test.iterate('content', COMPANY)] + [article for article in articles.iterate('content', COMPANY)])
y_predict = classifier.predict(X_predict)

# Count positive and negative classified articles
articles_positive = collections.Counter(np.array([dt.date() for dt in articles_test.iterate('datetime', COMPANY)] + [dt.date() for dt in articles.iterate('datetime', COMPANY)])[np.array(y_predict) == "positive"])
articles_negative = collections.Counter(np.array([dt.date() for dt in articles_test.iterate('datetime', COMPANY)] + [dt.date() for dt in articles.iterate('datetime', COMPANY)])[np.array(y_predict) == "negative"])
# Flip the sign of the negative values
articles_negative = dict([(key, -value) for key, value in articles_negative.items()])

print "Stock price datas: ", stock_prices.stock_data.index[0], stock_prices.stock_data.index[-1]
print "Article dates: ", articles_negative.keys()[0], articles_positive.keys()[-1]

# PLOT
BAR_WIDTH = 0.35
str_start_date = datetime.datetime.strptime(START_DATE,'%Y-%m-%d')
str_end_date = datetime.datetime.strptime(END_DATE,'%Y-%m-%d')

# Price
pl.figure(1)
pl.suptitle(COMPANY) # Main title
pl.subplot(311)
pl.plot_date(stock_prices.stock_data.index, stock_prices.stock_data[COMPANY], '-')
pl.title("Stock price")
pl.xlim(str_start_date, str_end_date)
#pl.xticks(rotation=30)
pl.tick_params(\
    axis='x',          # changes apply to the x-axis
    which='both',      # both major and minor ticks are affected
    bottom='off',      # ticks along the bottom edge are off
    top='off',         # ticks along the top edge are off
    labelbottom='off') # labels along the bottom edge are off
pl.grid(True)

# Average change bar chart
pl.subplot(312)
pl.title("Average stock change")
change = {}
for date in stock_prices.stock_data.index:
    change[date] = stock_prices.get_stock_change(COMPANY, date, AVERAGE_SPAN_DAYS) -1
pl.bar(change.keys(), change.values(), BAR_WIDTH, color='r')
pl.xlim(str_start_date, str_end_date)
#pl.xticks(rotation=30)
pl.tick_params(\
    axis='x',          # changes apply to the x-axis
    which='both',      # both major and minor ticks are affected
    bottom='off',      # ticks along the bottom edge are off
    top='off',         # ticks along the top edge are off
    labelbottom='off') # labels along the bottom edge are off
pl.grid(True)

# Article sentiment bar chart
pl.subplot(313)
pl.title("Article sentiment")
pl.bar(articles_negative.keys(), articles_negative.values(), BAR_WIDTH, color='r')
pl.bar(articles_positive.keys(), articles_positive.values(), BAR_WIDTH, color='g')
pl.xlim(str_start_date, str_end_date)
pl.xticks(rotation=30)
pl.grid(True)

pl.show()